// Copyright (c) 2009-2019 The Regents of the University of Michigan
// This file is part of the HOOMD-blue project, released under the BSD 3-Clause License.


// this include is necessary to get MPI included before anything else to support intel MPI
#include "hoomd/ExecutionConfiguration.h"

#include <functional>
#include <fstream>

#include "hoomd/md/TablePotential.h"
#include "hoomd/md/NeighborListTree.h"
#ifdef ENABLE_CUDA
#include "hoomd/md/TablePotentialGPU.h"
#endif

using namespace std;
using namespace std::placeholders;

#include "hoomd/test/upp11_config.h"

HOOMD_UP_MAIN();

/*! \file table_potential.cc
    \brief Implements unit tests for TablePotential and descendants
    \ingroup unit_tests
*/

//! Typedef'd TablePotential factory
typedef std::function<std::shared_ptr<TablePotential> (std::shared_ptr<SystemDefinition> sysdef,
                                                    std::shared_ptr<NeighborList> nlist,
                                                    unsigned int width)> table_potential_creator;

//! performs some really basic checks on the TablePotential class
void table_potential_basic_test(table_potential_creator table_creator, std::shared_ptr<ExecutionConfiguration> exec_conf)
    {
    // perform a basic test to see of the potential and force can be interpolated between two particles
    std::shared_ptr<SystemDefinition> sysdef_2(new SystemDefinition(2, BoxDim(1000.0), 1, 0, 0, 0, 0, exec_conf));
    std::shared_ptr<ParticleData> pdata_2 = sysdef_2->getParticleData();

    {
    ArrayHandle<Scalar4> h_pos(pdata_2->getPositions(), access_location::host, access_mode::readwrite);
    h_pos.data[0].x = h_pos.data[0].y = h_pos.data[0].z = 0.0;
    h_pos.data[1].x = Scalar(1.0); h_pos.data[1].y = h_pos.data[1].z = 0.0;
    }

    std::shared_ptr<NeighborListTree> nlist_2(new NeighborListTree(sysdef_2, Scalar(7.0), Scalar(0.8)));
    std::shared_ptr<TablePotential> fc_2 = table_creator(sysdef_2, nlist_2, 3);

    // first check for proper initialization by seeing if the force and potential come out to be 0
    fc_2->compute(0);

    {
    GlobalArray<Scalar4>& force_array_1 =  fc_2->getForceArray();
    GlobalArray<Scalar>& virial_array_1 =  fc_2->getVirialArray();
    unsigned int pitch = virial_array_1.getPitch();
    ArrayHandle<Scalar4> h_force_1(force_array_1,access_location::host,access_mode::read);
    ArrayHandle<Scalar> h_virial_1(virial_array_1,access_location::host,access_mode::read);
    MY_CHECK_SMALL(h_force_1.data[0].x, tol_small);
    MY_CHECK_SMALL(h_force_1.data[0].y, tol_small);
    MY_CHECK_SMALL(h_force_1.data[0].z, tol_small);
    MY_CHECK_SMALL(h_force_1.data[0].w, tol_small);
    MY_CHECK_SMALL(h_virial_1.data[0*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[1*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[2*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[3*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[4*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[5*pitch+0], tol_small);

    MY_CHECK_SMALL(h_force_1.data[1].x, tol_small);
    MY_CHECK_SMALL(h_force_1.data[1].y, tol_small);
    MY_CHECK_SMALL(h_force_1.data[1].z, tol_small);
    MY_CHECK_SMALL(h_force_1.data[1].w, tol_small);
    MY_CHECK_SMALL(h_virial_1.data[0*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[1*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[2*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[3*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[4*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_1.data[5*pitch+1], tol_small);
    }

    // specify a table to interpolate
    vector<Scalar> V, F;
    V.push_back(10.0);  F.push_back(1.0);
    V.push_back(21.0);  F.push_back(6.0);
    V.push_back(5.0);   F.push_back(2.0);
    fc_2->setTable(0, 0, V, F, 2.0, 4.0);

    // compute the forces again and check that they are still 0
    fc_2->compute(1);

    {
    GlobalArray<Scalar4>& force_array_2 =  fc_2->getForceArray();
    GlobalArray<Scalar>& virial_array_2 =  fc_2->getVirialArray();
    unsigned int pitch = virial_array_2.getPitch();
    ArrayHandle<Scalar4> h_force_2(force_array_2,access_location::host,access_mode::read);
    ArrayHandle<Scalar> h_virial_2(virial_array_2,access_location::host,access_mode::read);
    MY_CHECK_SMALL(h_force_2.data[0].x, tol_small);
    MY_CHECK_SMALL(h_force_2.data[0].y, tol_small);
    MY_CHECK_SMALL(h_force_2.data[0].z, tol_small);
    MY_CHECK_SMALL(h_force_2.data[0].w, tol_small);
    MY_CHECK_SMALL(h_virial_2.data[0*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[1*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[2*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[3*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[4*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[5*pitch+0], tol_small);

    MY_CHECK_SMALL(h_force_2.data[1].x, tol_small);
    MY_CHECK_SMALL(h_force_2.data[1].y, tol_small);
    MY_CHECK_SMALL(h_force_2.data[1].z, tol_small);
    MY_CHECK_SMALL(h_force_2.data[1].w, tol_small);
    MY_CHECK_SMALL(h_virial_2.data[0*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[1*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[2*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[3*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[4*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_2.data[5*pitch+1], tol_small);

    }

    // now go to rmin and check for the correct force value
    {
    ArrayHandle<Scalar4> h_pos(pdata_2->getPositions(), access_location::host, access_mode::readwrite);
    h_pos.data[1].x = Scalar(2.0);
    }

    fc_2->compute(2);

    {
    GlobalArray<Scalar4>& force_array_3 =  fc_2->getForceArray();
    GlobalArray<Scalar>& virial_array_3 =  fc_2->getVirialArray();
    unsigned int pitch = virial_array_3.getPitch();
    ArrayHandle<Scalar4> h_force_3(force_array_3,access_location::host,access_mode::read);
    ArrayHandle<Scalar> h_virial_3(virial_array_3,access_location::host,access_mode::read);
    MY_CHECK_CLOSE(h_force_3.data[0].x, -1.0, tol);
    MY_CHECK_SMALL(h_force_3.data[0].y, tol_small);
    MY_CHECK_SMALL(h_force_3.data[0].z, tol_small);
    MY_CHECK_CLOSE(h_force_3.data[0].w, 5.0, tol);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_3.data[0*pitch+0]
                                       +h_virial_3.data[3*pitch+0]
                                       +h_virial_3.data[5*pitch+0]), (1.0 / 6.0) * 2.0, tol);

    MY_CHECK_CLOSE(h_force_3.data[1].x, 1.0, tol);
    MY_CHECK_SMALL(h_force_3.data[1].y, tol_small);
    MY_CHECK_SMALL(h_force_3.data[1].z, tol_small);
    MY_CHECK_CLOSE(h_force_3.data[1].w, 5.0, tol);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_3.data[0*pitch+1]
                                       +h_virial_3.data[3*pitch+1]
                                       +h_virial_3.data[5*pitch+1]), (1.0 / 6.0) * 2.0, tol);
    }

    // go halfway in-between two points
    {
    ArrayHandle<Scalar4> h_pos(pdata_2->getPositions(), access_location::host, access_mode::readwrite);
    h_pos.data[1].y = Scalar(3.5);
    h_pos.data[1].x = Scalar(0.0);
    }

    // check the forces
    fc_2->compute(3);

    {
    GlobalArray<Scalar4>& force_array_4 =  fc_2->getForceArray();
    GlobalArray<Scalar>& virial_array_4 =  fc_2->getVirialArray();
    unsigned int pitch = virial_array_4.getPitch();
    ArrayHandle<Scalar4> h_force_4(force_array_4,access_location::host,access_mode::read);
    ArrayHandle<Scalar> h_virial_4(virial_array_4,access_location::host,access_mode::read);
    MY_CHECK_CLOSE(h_force_4.data[0].y, -4.0, tol);
    MY_CHECK_SMALL(h_force_4.data[0].x, tol_small);
    MY_CHECK_SMALL(h_force_4.data[0].z, tol_small);
    MY_CHECK_CLOSE(h_force_4.data[0].w, 13.0/2.0, tol);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_4.data[0*pitch+0]
                                       +h_virial_4.data[3*pitch+0]
                                       +h_virial_4.data[5*pitch+0]), (1.0 / 6.0) * 4.0 * 3.5, tol);

    MY_CHECK_CLOSE(h_force_4.data[1].y, 4.0, tol);
    MY_CHECK_SMALL(h_force_4.data[1].x, tol_small);
    MY_CHECK_SMALL(h_force_4.data[1].z, tol_small);
    MY_CHECK_CLOSE(h_force_4.data[1].w, 13.0 / 2.0, tol);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_4.data[0*pitch+1]
                                       +h_virial_4.data[3*pitch+1]
                                       +h_virial_4.data[5*pitch+1]), (1.0 / 6.0) * 4.0 * 3.5, tol);
    }

    // and now check for when r > rmax
    {
    ArrayHandle<Scalar4> h_pos(pdata_2->getPositions(), access_location::host, access_mode::readwrite);
    h_pos.data[1].z = Scalar(4.0);
    }

    // compute and check
    fc_2->compute(4);

    {
    GlobalArray<Scalar4>& force_array_5 =  fc_2->getForceArray();
    GlobalArray<Scalar>& virial_array_5 =  fc_2->getVirialArray();
    unsigned int pitch = virial_array_5.getPitch();
    ArrayHandle<Scalar4> h_force_5(force_array_5,access_location::host,access_mode::read);
    ArrayHandle<Scalar> h_virial_5(virial_array_5,access_location::host,access_mode::read);
    MY_CHECK_SMALL(h_force_5.data[0].x, tol_small);
    MY_CHECK_SMALL(h_force_5.data[0].y, tol_small);
    MY_CHECK_SMALL(h_force_5.data[0].z, tol_small);
    MY_CHECK_SMALL(h_force_5.data[0].w, tol_small);
    MY_CHECK_SMALL(h_virial_5.data[0*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[1*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[2*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[3*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[4*pitch+0], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[5*pitch+0], tol_small);

    MY_CHECK_SMALL(h_force_5.data[1].x, tol_small);
    MY_CHECK_SMALL(h_force_5.data[1].y, tol_small);
    MY_CHECK_SMALL(h_force_5.data[1].z, tol_small);
    MY_CHECK_SMALL(h_force_5.data[1].w, tol_small);
    MY_CHECK_SMALL(h_virial_5.data[0*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[1*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[2*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[3*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[4*pitch+1], tol_small);
    MY_CHECK_SMALL(h_virial_5.data[5*pitch+1], tol_small);
    }
    }

//! checks to see if TablePotential correctly handles multiple types
void table_potential_type_test(table_potential_creator table_creator, std::shared_ptr<ExecutionConfiguration> exec_conf)
    {
    // perform a basic test to see of the potential and force can be interpolated between two particles
    std::shared_ptr<SystemDefinition> sysdef(new SystemDefinition(4, BoxDim(1000.0), 2, 0, 0, 0, 0, exec_conf));
    std::shared_ptr<ParticleData> pdata = sysdef->getParticleData();

    {
    ArrayHandle<Scalar4> h_pos(pdata->getPositions(), access_location::host, access_mode::readwrite);
    h_pos.data[0].x = h_pos.data[0].y = h_pos.data[0].z = 0.0; h_pos.data[0].w= __int_as_scalar(0);
    h_pos.data[1].x = Scalar(1.5); h_pos.data[1].y = h_pos.data[1].z = 0.0; h_pos.data[1].w = __int_as_scalar(1);
    h_pos.data[2].x = 0.0; h_pos.data[2].y = Scalar(1.5); h_pos.data[2].z = 0.0; h_pos.data[2].w = __int_as_scalar(0);
    h_pos.data[3].x = Scalar(1.5); h_pos.data[3].y = Scalar(1.5); h_pos.data[3].z = 0.0; h_pos.data[3].w = __int_as_scalar(1);
    }

    std::shared_ptr<NeighborListTree> nlist(new NeighborListTree(sysdef, Scalar(2.0), Scalar(0.8)));
    std::shared_ptr<TablePotential> fc = table_creator(sysdef, nlist, 3);

    // specify a table to interpolate
    vector<Scalar> V, F;
    V.push_back(10.0);  F.push_back(1.0);
    V.push_back(20.0);  F.push_back(6.0);
    V.push_back(5.0);   F.push_back(2.0);
    fc->setTable(0, 0, V, F, 1.0, 2.0);

    // next type pair
    V.clear(); F.clear();
    V.push_back(20.0);  F.push_back(2.0);
    V.push_back(40.0);  F.push_back(12.0);
    V.push_back(10.0);   F.push_back(4.0);
    fc->setTable(0, 1, V, F, 0.0, 2.0);

    // next type pair
    V.clear(); F.clear();
    V.push_back(5.0);  F.push_back(0.5);
    V.push_back(10.0);  F.push_back(3.0);
    V.push_back(2.5);   F.push_back(2.0);
    fc->setTable(1, 1, V, F, 1.0, 2.0);

    // compute and check
    fc->compute(0);

    {
    GlobalArray<Scalar4>& force_array_6 =  fc->getForceArray();
    GlobalArray<Scalar>& virial_array_6 =  fc->getVirialArray();
    unsigned int pitch = virial_array_6.getPitch();
    ArrayHandle<Scalar4> h_force_6(force_array_6,access_location::host,access_mode::read);
    ArrayHandle<Scalar> h_virial_6(virial_array_6,access_location::host,access_mode::read);
    MY_CHECK_CLOSE(h_force_6.data[0].x, -8.0, tol);
    MY_CHECK_CLOSE(h_force_6.data[0].y, -6.0, tol);
    MY_CHECK_SMALL(h_force_6.data[0].z, tol_small);
    MY_CHECK_SMALL(h_force_6.data[0].w, 10.0+25.0);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_6.data[0*pitch+0]
                                       +h_virial_6.data[3*pitch+0]
                                       +h_virial_6.data[5*pitch+0]), (8*1.5+6*1.5)*1.0/6.0, tol);

    MY_CHECK_CLOSE(h_force_6.data[1].x, 8.0, tol);
    MY_CHECK_CLOSE(h_force_6.data[1].y, -3.0, tol);
    MY_CHECK_SMALL(h_force_6.data[1].z, tol_small);
    MY_CHECK_CLOSE(h_force_6.data[1].w, 25.0/2.0 + 5.0, tol);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_6.data[0*pitch+1]
                                       +h_virial_6.data[3*pitch+1]
                                       +h_virial_6.data[5*pitch+1]), (8*1.5 + 3.0 * 1.5)*1.0/6.0, tol);

    MY_CHECK_CLOSE(h_force_6.data[2].x, -8.0, tol);
    MY_CHECK_CLOSE(h_force_6.data[2].y, 6.0, tol);
    MY_CHECK_SMALL(h_force_6.data[2].z, tol_small);
    MY_CHECK_SMALL(h_force_6.data[2].w, 10.0+25.0);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_6.data[0*pitch+2]
                                       +h_virial_6.data[3*pitch+2]
                                       +h_virial_6.data[5*pitch+2]), (8*1.5+6*1.5)*1.0/6.0, tol);

    MY_CHECK_CLOSE(h_force_6.data[3].x, 8.0, tol);
    MY_CHECK_CLOSE(h_force_6.data[3].y, 3.0, tol);
    MY_CHECK_SMALL(h_force_6.data[3].z, tol_small);
    MY_CHECK_CLOSE(h_force_6.data[3].w, 25.0/2.0 + 5.0, tol);
    MY_CHECK_CLOSE(Scalar(1./3.)*(h_virial_6.data[0*pitch+3]
                                       +h_virial_6.data[3*pitch+3]
                                       +h_virial_6.data[5*pitch+3]), (8*1.5 + 3.0*1.5)*1.0/6.0, tol);
    }
    }

//! TablePotential creator for unit tests
std::shared_ptr<TablePotential> base_class_table_creator(std::shared_ptr<SystemDefinition> sysdef,
                                                    std::shared_ptr<NeighborList> nlist,
                                                    unsigned int width)
    {
    return std::shared_ptr<TablePotential>(new TablePotential(sysdef, nlist, width));
    }

#ifdef ENABLE_CUDA
//! TablePotentialGPU creator for unit tests
std::shared_ptr<TablePotential> gpu_table_creator(std::shared_ptr<SystemDefinition> sysdef,
                                             std::shared_ptr<NeighborList> nlist,
                                             unsigned int width)
    {
    nlist->setStorageMode(NeighborList::full);
    std::shared_ptr<TablePotentialGPU> table(new TablePotentialGPU(sysdef, nlist, width));
    return table;
    }
#endif


//! test case for basic test on CPU
UP_TEST( TablePotential_basic )
    {
    table_potential_creator table_creator_base = bind(base_class_table_creator, _1, _2, _3);
    table_potential_basic_test(table_creator_base, std::shared_ptr<ExecutionConfiguration>(new ExecutionConfiguration(ExecutionConfiguration::CPU)));
    }

//! test case for type test on CPU
UP_TEST( TablePotential_type )
    {
    table_potential_creator table_creator_base = bind(base_class_table_creator, _1, _2, _3);
    table_potential_type_test(table_creator_base, std::shared_ptr<ExecutionConfiguration>(new ExecutionConfiguration(ExecutionConfiguration::CPU)));
    }

#ifdef ENABLE_CUDA
//! test case for basic test on GPU
UP_TEST( TablePotentialGPU_basic )
    {
    table_potential_creator table_creator_gpu = bind(gpu_table_creator, _1, _2, _3);
    table_potential_basic_test(table_creator_gpu, std::shared_ptr<ExecutionConfiguration>(new ExecutionConfiguration(ExecutionConfiguration::GPU)));
    }

//! test case for type test on GPU
UP_TEST( TablePotentialGPU_type )
    {
    table_potential_creator table_creator_gpu = bind(gpu_table_creator, _1, _2, _3);
    table_potential_type_test(table_creator_gpu, std::shared_ptr<ExecutionConfiguration>(new ExecutionConfiguration(ExecutionConfiguration::GPU)));
    }
#endif
